# -*- coding: utf-8 -*-
"""
Created on Wed May 21 11:57:51 2025

@author: lenovo
"""

import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime
import glob
from tqdm import tqdm
# 自定义label
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
from gsj_typhoon import split_str_id,tydat_NH,tydat


show_fig = False
names,track_ids = ['2023NH42_36','2023NH47_46'],[42,47]
tynames,tyids = split_str_id(names)

for idx,(tyname,tyid) in enumerate(zip(tynames,tyids)) :
        
    RIstd = 7
    path= rf'D:\met_data\ty_ensemble\{tyname}_{tyid}'
    dates=os.listdir(path)
    paths=[path+"\\"+n for n in dates]
    
    pic_savepath = rf'D:\met_data\pics\{names[idx]}-RI{RIstd}\w_p_track'
    os.makedirs(pic_savepath,exist_ok=True)
    
    
    def basemap(ax,ty,title=None):
        ax.add_feature(cfeature.COASTLINE,linewidth=0.5) #海岸线
        ax.add_feature(cfeature.BORDERS, linestyle=':')   #国界
        china_provinces = cfeature.NaturalEarthFeature(category='cultural',name='admin_1_states_provinces_lines', scale='10m', facecolor='none')
        ax.add_feature(china_provinces, edgecolor='black', linewidth=0.5) #省界
        ax.set(xlim=(ty.extent[0],ty.extent[1]))
        ax.set(ylim=(ty.extent[2],ty.extent[3]))    
        #这种设定方法适用于这样子给出范围。
        ax.set_xticks(np.linspace(ty.extent[0], ty.extent[1]+5, 5))
        ax.set_yticks(np.linspace(ty.extent[2], ty.extent[3]+5, 5))
        ax.xaxis.set_major_formatter(LongitudeFormatter(number_format='.1f'))
        ax.xaxis.set_minor_locator(plt.MultipleLocator(1))
        ax.yaxis.set_major_formatter(LatitudeFormatter(number_format='.1f'))
        ax.yaxis.set_minor_locator(plt.MultipleLocator(1))
        ax.tick_params(axis='both', labelsize=6, direction='out')
        ax.set_title(title)
    
    def track(ax,ty,sep=1,color='Black'):
        lon= ty.lon
        lat= ty.lat
        ll2=np.arange(0,len(lon)-1,1)
        # 连线
        for i in ll2:
            pointA = lon[i],lat[i]
            pointB = lon[i+1],lat[i+1]
            ax.add_geometries([sgeom.LineString([pointA, pointB])], color=color,crs=ccrs.PlateCarree())  
        # legend
        legend_elements = [
            Line2D([0], [0], color='black', lw=1, label='emsemble members',marker='o'),
            Line2D([0], [0], color='green', lw=1, label='mean of members'),
            Line2D([0], [0], color='red', lw=1, label='ERA5'),
            Line2D([0], [0], color='blue', lw=1, label='observation'),
            ]
        ax.legend( handles=legend_elements,loc='lower right',fontsize=8 )
    
    # Line2D([0], [0], color='yellow', lw=1, marker='o', label='Marker only',ls='none',markersize=5)
    
    def SLP(ax,ty,title,color="Black"):
        ax.plot(ty.time, ty.pmin, 'o-', color=color)
        # legend
        legend_elements = [
            Line2D([0], [0], color='black', lw=1, label='emsemble members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='green', lw=1, label='mean of members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='red', lw=1, label='ERA5',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='blue', lw=1, label='observation',marker='o',markersize=5,ls="none"),
            ]
        ax.legend( handles=legend_elements,loc='lower right',fontsize=10 )    
        # label
        min_slp_label = np.arange(920,1030,10)
        ax.set_title(title)
        ax.set_xlabel('Time (month-day)')
        ax.set_ylabel('min SLP (hPa)')
        ax.set_yticks( min_slp_label )
        ax.tick_params(axis='y')
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=2))  # 每5天一个主刻度
        ax.xaxis.set_minor_locator(mdates.DayLocator(interval=1))  # 每1天一个次刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d"))  # 主刻度显示的格式
    
    
    def maxw(ax,ty,title,color="Black"):
        ax.plot(ty.time, ty.umax, 'o-', color=color)
        # legend
        legend_elements = [
            Line2D([0], [0], color='black', lw=1, label='emsemble members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='green', lw=1, label='mean of members',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='red', lw=1, label='ERA5',marker='o',markersize=5,ls="none"),
            Line2D([0], [0], color='blue', lw=1, label='observation',marker='o',markersize=5,ls="none"),
            ]
        ax.legend( handles=legend_elements,loc='upper right',fontsize=10 )   
        # label
        max_wind_label = np.arange(0,70,10)
        ax.set_title(title)
        ax.set_xlabel('Time (month-day)')
        ax.set_ylabel('max wind (m/s)')
        ax.set_yticks( max_wind_label )
        ax.tick_params(axis='y')
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=2))  # 每5天一个主刻度
        ax.xaxis.set_minor_locator(mdates.DayLocator(interval=1))  # 每1天一个次刻度
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%m-%d"))  # 主刻度显示的格式
        
        
    
    for i in tqdm(range(len(paths)),desc = f'processing {tyname}'):
        try:
            path=paths[i]
            title=dates[i]+tyname
            # mean
            meanpath = glob.glob(os.path.join(path,f"*{tyid}_mean"))[0]
            # reanalysis
            n0path= path+r"\TRACK_ID_0"
            
            # other path
            start = 'TRACK_ID'   # 获取目录下所有文件和子目录
            all_files = os.listdir(path)   # 过滤出文件且符合后缀要求
            filtered_files = [f for f in all_files if f.startswith(start) and f != "TRACK_ID_0"]   # 打印所有符合后缀的文件名
            
        ##################  track #################
            fig,ax= plt.subplots(subplot_kw={'projection': ccrs.PlateCarree()},dpi=300,figsize=(10,5))
            #绘制track 图
            for f in filtered_files:
                ty=tydat(path+"\\"+f,7)
                track(ax,ty,color="Black") if ty.control.any() else None
            # 绘制mean
            ty=tydat(meanpath,7,skiprows=5)  
            basemap(ax,ty,title=title)  if ty.control.any() else None
            track(ax,ty,color="lime")
            # 绘制trackid 0
            ty=tydat(n0path,7,skiprows=2)  
            track(ax,ty,color="Red")
            # 绘制实况
            tyobs = tydat_NH(track_id=track_ids[idx])
            track(ax,tyobs,color="blue")
            if not show_fig:
                plt.savefig(f"{pic_savepath}\\track_{title}")
                plt.close()
            else:
                plt.show()
            
        ################## SLP  #################    
            fig,ax= plt.subplots(dpi=300,figsize=(10,5))
            #绘制集合成员
            
            for f in filtered_files:
                ty=tydat(path+"\\"+f,7)
                SLP(ax,ty,title,color="Black")
            # 绘制mean
            ty=tydat(meanpath,7,skiprows=5)  
            SLP(ax,ty,title,color="lime") 
            # 绘制trackid 0
            ty=tydat(n0path,7,skiprows=2)  
            SLP(ax,ty,title,color="Red")
            # 绘制 实况
            SLP(ax,tyobs,title,color="blue")
            if show_fig:
                plt.show()
            else:
                plt.savefig(f"{pic_savepath}\\SLP_{title}")
                plt.close()
        
        #################  maxw   #################
            fig,ax= plt.subplots(dpi=300,figsize=(10,5));
            #绘制集合成员
            for f in filtered_files:
                ty=tydat(path+"\\"+f,7)
                maxw(ax,ty,title,color="Black")
            # 绘制mean
            ty=tydat(meanpath,7,skiprows=5)  
            maxw(ax,ty,title,color="lime") 
            # 绘制trackid 0
            ty=tydat(n0path,7,skiprows=2)  
            maxw(ax,ty,title,color="Red")
            #  绘制实况
            maxw(ax,tyobs,title,color="blue")
            if not show_fig:
                plt.savefig(f"{pic_savepath}\\maxw_{title}")
                plt.close()
            else:
                plt.show()
        except Exception as e:
            print(f'{path}存在缺测情况,报错\n{e}')